#include <stdint.h>   // for uintptr_t

#define BALLOC_SECURE true
#define BALLOC_IMPLEMENTATION   // implementation file
#undef BALLOC_ALIGNMENT        // change pointer alignment for testing
#define BALLOC_ALIGNMENT 8
#include "balloc.h"

#include "acutest.h"

#ifndef UNUSED
#define UNUSED(X) ((void)X)
#endif

static int errors_logged;

static void test_error_handler(const char *file, int line, const char *msg)
{
    UNUSED(file); UNUSED(line); UNUSED(msg);
    errors_logged++;
    TEST_MSG("Error triggered: %s(%d) %s", file, line, msg);
}

void test_dynamic_pool(void)
{
    balloc_error_handler = test_error_handler;
    errors_logged = 0;

    balloc_arena arena = balloc_new_arena(32);
    TEST_CHECK(arena != NULL);
    balloc_arena_stats arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 0);
    TEST_CHECK(arena_stats.error_count == 0);

    void *p1 = balloc(arena, 8);
    TEST_CHECK(p1 != NULL);
    TEST_CHECK(((uintptr_t)p1 % BALLOC_ALIGNMENT) == 0);
    arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 8);
    TEST_CHECK(arena_stats.error_count == 0);

    void *p2 = balloc(arena, 16);
    TEST_CHECK(p2 != NULL);
    arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 24);
    TEST_CHECK(arena_stats.error_count == 0);

    // Allocation beyond capacity should fail
    void *p3 = balloc(arena, 16);
    TEST_CHECK(p3 == NULL);
    TEST_CHECK(errors_logged == 1);
    arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 24);
    TEST_CHECK(arena_stats.error_count == 1);

    balloc_free(arena);
}

void test_static_pool(void)
{
    balloc_error_handler = test_error_handler;
    errors_logged = 0;

    char pool[32];
    balloc_arena arena = balloc_new_static_arena(pool, sizeof(pool));
    TEST_CHECK(arena != NULL);
    balloc_arena_stats arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 0);
    TEST_CHECK(arena_stats.error_count == 0);

    void *p1 = balloc(arena, 8);
    TEST_CHECK(p1 != NULL);
    TEST_CHECK(((uintptr_t)p1 % BALLOC_ALIGNMENT) == 0);
    arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 8);
    TEST_CHECK(arena_stats.error_count == 0);

    void *p2 = balloc(arena, 16);
    TEST_CHECK(p2 != NULL);
    arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 24);
    TEST_CHECK(arena_stats.error_count == 0);

    // Allocation beyond capacity should fail
    void *p3 = balloc(arena, 16);
    TEST_CHECK(p3 == NULL);
    TEST_CHECK(errors_logged == 1);
    arena_stats = balloc_get_arena_stats(arena);
    TEST_CHECK(arena_stats.current_size == 32);
    TEST_CHECK(arena_stats.current_used == 24);
    TEST_CHECK(arena_stats.error_count == 1);

    balloc_free(arena);
}

void test_alignment(void)
{
    balloc_arena arena = balloc_new_arena(64);
    TEST_CHECK(arena != NULL);

    for (int i = 0; i < 5; ++i)
    {
        void *p = balloc(arena, 5);  // Request intentionally misaligned size
        TEST_CHECK(p != NULL);
        TEST_CHECK(((uintptr_t)p % BALLOC_ALIGNMENT) == 0);
    }

    balloc_free(arena);
}

TEST_LIST =
{
    { "dynamic_pool", test_dynamic_pool },
    { "static_pool", test_static_pool },
    { "alignment", test_alignment },
    { NULL, NULL }
};
